import os


anno_path = '/home/sean/data/UWdevkit/Annotations/train/'
train_video_file = '/home/sean/data/UWdevkit/ImageSets/train_video.txt'
train_file = '/home/sean/data/UWdevkit/ImageSets/train.txt'
test_file = '/home/sean/data/UWdevkit/ImageSets/test.txt'

exclusive = ['n01904806', 'n09432283', 'qdvt', 'trepang07', 'urchin_baidu', 'urchin_google',
             'vs170622', 'zhongkeyuan3', 'q3_RUAS', 'q2_RUAS', 'q8_RUAS', 'qdtv', 'shell3',
             'shell6', 'trepang6', 'val_08']
# print(no_object_lines)
## Video
# with open(train_video_file, 'w') as wf:
#     for dir in os.listdir(anno_path):
#         if dir not in exclusive:
#             sub_dir = os.path.join(anno_path, dir)
#             file = os.listdir(sub_dir)
#             write_path = dir +' ' +str(len(file)) +'\n'
#             print(write_path)
#             wf.write(write_path)

# SSD train
with open(train_file, 'w') as wf:
    for dir in os.listdir(anno_path):
        if dir.split('_')[0] != 'val':
            sub_dir = os.path.join(anno_path, dir)
            for file in os.listdir(sub_dir):
                write_path = os.path.join(dir, file.split('.')[0]) + '\n'
                print(write_path)
                wf.write(write_path)

## test

# with open(test_file, 'w') as wf:
#     for dir in os.listdir(anno_path):
#         sub_dir = os.path.join(anno_path, dir)
#         files = os.listdir(sub_dir)
#         files_num = [f.split('.')[0] for f in files]
#         sort_files = sorted(files_num)
#         for file in sort_files:
#             # print(file)
#             write_path = os.path.join(dir, file) + '\n'
#             print(write_path)
#             wf.write(write_path)
